#!/usr/bin/env python

DBG=False

if DBG:
  from hexdump import hexdump
  from textcolor import *

  # kill all children on error
  os.setpgrp() # create new process group, become its leader
  assert os.getpgrp() == os.getpid()
  os._exit = lambda _: os.killpg(0, signal.SIGKILL)
  import atexit
  atexit.register(lambda: os.killpg(0, signal.SIGKILL))

import os,sys,struct,signal,re
from errno import *
from ctypes import *
from ctypes.util import find_library
libc = cdll.LoadLibrary(find_library('c'))
_ptrace = libc.ptrace
_ptrace.argtypes, _ptrace.restype = (c_ulong, c_ulong, c_ulong, c_ulong), c_ulong
CPU_WORD_SIZE = sizeof(c_void_p)
TRACEME,PEEKDATA,POKEDATA,GETREGS,SETREGS,SYSCALL,SETOPTIONS,GETEVENTMSG =0,2,5,12,13,24,0x4200,0x4201
TRACESYSGOOD,TRACEFORK,TRACEVFORK,TRACECLONE,TRACEEXEC,TRACEVFORKDONE,TRACEEXIT =1,2,4,8,16,32,64
EVENT_FORK,EVENT_VFORK,EVENT_CLONE,EVENT_EXEC,EVENT_VFORK_DONE,EVENT_EXIT =1,2,3,4,5,6

def anyfile(names):
 for n in names:
   try: return file(n)
   except: pass
 raise IOError, 'found none of files %s' % names

# get syscall numbers
SYSCALL_NAMES = {}
for line in anyfile({4:['/usr/include/asm/unistd_32.h'],
                     8:['/usr/include/asm/unistd_64.h','/usr/include/asm-x86_64/unistd.h']}[CPU_WORD_SIZE]):
  r = re.search('^#define\s+__NR_(\S+)\s+(\d+)', line)
  if r:
    name, num = r.group(1), int(r.group(2))
    SYSCALL_NAMES[num] = name

def ptrace(command, pid=0, arg1=0, arg2=0):
  result = _ptrace(command, pid, arg1, arg2)
  assert result!=-1, "error ptrace(cmd=%s, pid=%s, %r, %r)" % (command, pid, arg1, arg2)
  return result

class ptrace_registers_t(Structure):
  if CPU_WORD_SIZE == 4:
    _fields_ = [(n, c_long) for n in 'ebx ecx edx esi edi ebp acc ds es fs gs orig_acc eip cs eflags esp ss'.split()]
    def __getitem__(self, i):
      return (self.ebx, self.ecx, self.edx, self.esi, self.edi, self.ebp)[i]
    def __setitem__(self, i, v):
      if i==0: self.ebx=v
      if i==1: self.rcx=v
      if i==2: self.edx=v
      if i==3: self.esi=v
      if i==4: self.edi=v
      if i==5: self.ebp=v
  elif CPU_WORD_SIZE == 8:
    _fields_ = [(n, c_long) for n in 'r15 r14 r13 r12 rbp rbx r11 r10 r9 r8 acc rcx rdx rsi rdi orig_acc rip cs eflags rsp ss fs_base gs_base ds es fs gs'.split()]
    def __getitem__(self, i):
      return (self.rdi, self.rsi, self.rdx, self.r10, self.r8, self.r9)[i]
    def __setitem__(self, i, v):
      if i==0: self.rdi=v
      if i==1: self.rsi=v
      if i==2: self.rdx=v
      if i==3: self.r10=v
      if i==4: self.r8=v
      if i==5: self.r9=v
  # kludge for py2.4, with py2.6 copy.copy works
  def copy(self):
    p = ptrace_registers_t()
    memmove(addressof(p), addressof(self), sizeof(self))
    return p

class stat_t(Structure):
  if CPU_WORD_SIZE==4: # stat64
    _fields_ = [
      ('st_dev',          c_ulonglong),
      ('__pad0',          c_uint),
      ('__st_ino',        c_ulong),
      ('st_mode',         c_uint),
      ('st_nlink',        c_uint),
      ('st_uid',          c_ulong),
      ('st_gid',          c_ulong),
      ('st_rdev',         c_ulonglong),
      ('__pad3',          c_uint),
      ('st_size',         c_longlong),
      ('st_blksize',      c_ulong),
      ('st_blocks',       c_ulonglong),
      ('st_atime',        c_ulong), ('st_atime_nsec',   c_ulong),  # time of last access
      ('st_mtime',        c_ulong), ('st_mtime_nsec',   c_ulong),  # time of last modification
      ('st_ctime',        c_ulong), ('st_ctime_nsec',   c_ulong),  # time of last status change
      ('st_ino',          c_ulonglong)]
  elif CPU_WORD_SIZE==8:
    _fields_ = [
      ('st_dev',          c_ulong),   # ID of device containing file
      ('st_ino',          c_ulong),   # inode number
      ('st_nlink',        c_ulong),   # number of hard links
      ('st_mode',         c_uint),    # protection
      ('st_uid',          c_uint),    # user ID of owner
      ('st_gid',          c_uint),    # group ID of owner
      ('__pad0',          c_uint),
      ('st_rdev',         c_ulong),   # device ID (if special file)
      ('st_size',         c_long),    # total size, in bytes
      ('st_blksize',      c_long),    # blocksize for filesystem I/O
      ('st_blocks',       c_long),    # number of blocks allocated
      ('st_atime',        c_ulong), ('st_atime_nsec',   c_ulong),  # time of last access
      ('st_mtime',        c_ulong), ('st_mtime_nsec',   c_ulong),  # time of last modification
      ('st_ctime',        c_ulong), ('st_ctime_nsec',   c_ulong),  # time of last status change
      ('__unused',        c_long * 3)]

def alignup(x):
  return (x + CPU_WORD_SIZE - 1) & ~(CPU_WORD_SIZE - 1)

def designed(x):
  return x % 256**CPU_WORD_SIZE

def read_bytes(pid, address, length):
  fd = os.open("/proc/%d/mem" % pid, os.O_RDONLY | os.O_LARGEFILE)
  assert os.lseek(fd, designed(address), 0) == designed(address)
  s = os.read(fd, length)
  os.close(fd)
  return s

def read_string(pid, address, maxlength=0x200):
  fd = os.open("/proc/%d/mem" % pid, os.O_RDONLY | os.O_LARGEFILE)
  assert os.lseek(fd, designed(address), 0) == designed(address)
  s = ''
  while len(s)<maxlength:
    c = os.read(fd, 1)
    if c=='\0': break
    s += c
  os.close(fd)
  return s

def _write_word(pid, address, bytes):
  mask = {4:'=L', 8:'=Q'}[CPU_WORD_SIZE]
  assert address // CPU_WORD_SIZE == (address+len(bytes)-1) // CPU_WORD_SIZE
  if len(bytes) == CPU_WORD_SIZE:
    assert address % CPU_WORD_SIZE == 0
    ptrace(POKEDATA, pid, address, struct.unpack(mask, bytes)[0])
  else:
    bytes8 = struct.pack(mask, ptrace(PEEKDATA, pid, address))
    bytes8 = bytes8[:address % CPU_WORD_SIZE] + bytes + bytes8[(address % CPU_WORD_SIZE) + len(bytes):]
    ptrace(POKEDATA, pid, address, struct.unpack(mask, bytes8)[0])

def write_bytes(pid, address, bytes):
  while bytes:
    hdrsize = CPU_WORD_SIZE - (address % CPU_WORD_SIZE)
    _write_word(pid, address, bytes[:hdrsize])
    address += hdrsize
    bytes = bytes[hdrsize:]
  assert read_bytes(pid, address, len(bytes)) == bytes

# TODO: ttl
FILES = {os.path.abspath('setup.py'):  'hello!' * 100,
         os.path.abspath('1.so'):      'soso',
         os.path.abspath('1.py'):      'print file("1.so").read()',
         os.path.abspath('a.out'):     file('/root/a.out').read(),
         os.path.abspath('distorm.so'):file('/root/distorm.so').read(),
         os.path.abspath('2.py'):      ('from distorm import Decode, Decode32Bits\n'
                                        'print Decode(0x400000, "b800000000".decode("hex"), Decode32Bits)\n'),
         os.path.abspath('1.c'):       ('#include <stdio.h>\n'
                                        'int main() {\n'
                                        ' printf("%s\\n", "Hello, world");\n'
                                        ' return 0;\n'
                                        '}\n')
         }

if sys.argv[-1] != '--magic':
  pid = os.fork()
  if pid != 0: # parent
    HANDLES = {} # fd -> {name:, file:, pos:}

    def seek(fd, off, whence):
      H = HANDLES[fd]
      if whence==1: H['pos'] += off
      elif whence==2: H['pos'] = len(H['file']) + off
      else: H['pos'] = off
      H['pos'] = min(max(H['pos'], 0), len(H['file']))
      return H['pos']

    P = {pid:{}}
    pid, status = os.waitpid(pid, 0)
    ptrace(SETOPTIONS, pid, 0, TRACEFORK|TRACEVFORK|TRACECLONE)
    ptrace(SYSCALL, pid)
    while 1:
      try:
        pid, status = os.waitpid(-1, 0)
      except OSError, (errno, _):
        if errno==ECHILD: sys.exit() # No child processes
        raise

      if os.WIFEXITED(status):
        if DBG: colorwrite( '%d '%pid, WHITE)
        if DBG: colorwrite( 'exited', YELLOW)
        del P[pid]
        continue

      elif os.WIFSIGNALED(status):
        if DBG: print 'WIFSIGNALED %d %x' % (pid, status)
        continue

      elif not os.WIFSTOPPED(status):
        assert False, (pid, hex(status))

      elif (status >> 16) in (EVENT_FORK,EVENT_VFORK,EVENT_CLONE):
        new_pid = c_int()
        ptrace(GETEVENTMSG, pid, 0, addressof(new_pid))
        P[new_pid.value] = {'args':True}
        if DBG: print 'status', hex(status), 'newpid', new_pid.value
      elif (status >> 16) != 0:
        assert False, 'status %x' % status
      elif os.WSTOPSIG(status) != signal.SIGTRAP:
        if DBG: print hex(status), os.WIFEXITED(status), os.WIFSIGNALED(status), os.WIFSTOPPED(status)
        if os.WSTOPSIG(status) == signal.SIGSEGV: sys.exit()

      elif 'args' not in P[pid]: # === on before ===
        r = ptrace_registers_t()
        ptrace(GETREGS, pid, 0, addressof(r))

        P[pid]['args'] = r.copy()
        P[pid]['syscall'] = syscall = SYSCALL_NAMES.get(r.orig_acc, 'syscall_%d' % r.orig_acc)
        P[pid]['return'] = None

        if DBG: colorwrite( '%d '%pid, WHITE)
        if DBG: colorwrite( syscall + '(%x,%x,%x,%x,%x,%x)' % r[0:6], DARK_MAGENTA )

        # argument as path
        def p(n):
          return os.path.normpath(os.path.join(os.readlink('/proc/%d/cwd' % pid), read_string(pid, r[n])))

        def skipsyscall(acc):
          P[pid]['return'] = r.copy()
          P[pid]['return'].acc = acc
          assert r.orig_acc == P[pid]['return'].orig_acc, (r.orig_acc, P[pid]['return'].orig_acc)
          r.orig_acc += 9999  # avoid syscall by making its number invalid
          assert r.orig_acc != P[pid]['return'].orig_acc
          ptrace(SETREGS, pid, 0, addressof(r))

        def fakestat(f):
          st = stat_t()
          st.st_mtime = st.st_atime = st.st_ctime = 0x40000000 # fake time '10 Jan 2004'
          st.st_mode = 0100555 # for dir 0040555
          st.st_size = len(f)
          st.st_dev = 1
          write_bytes(pid, r[1], buffer(st))
          skipsyscall(0)

        if syscall=='open':
          name = p(0)
          if DBG: colorwrite( name, YELLOW)
          if name in FILES:
            if r[1] & 3:
              skipsyscall(-EROFS)
            else:
              fd = 10000
              if HANDLES: fd = max(HANDLES)+1
              HANDLES[fd] = {'name':name, 'file':FILES[name], 'pos':0}
              skipsyscall(fd)

        elif syscall in ('getxattr', 'lgetxattr', 'setxattr', 'lsetxattr', 'listxattr', 'llistxattr', 'removexattr', 'lremovexattr') and p(0) in FILES or\
             syscall in ('fgetxattr', 'fsetxattr', 'flistxattr', 'fremovexattr') and r[0] in HANDLES:
          skipsyscall(-EOPNOTSUPP)

        elif syscall=='newfstatat':
          name = p(1)
          if DBG: colorwrite( name, YELLOW)
          if name in FILES: fakestat(FILES[name])

        elif CPU_WORD_SIZE==8 and syscall in ('stat', 'lstat') or\
             CPU_WORD_SIZE==4 and syscall in ('stat64', 'lstat64'):
          name = p(0)
          if DBG: colorwrite( name, YELLOW)
          if name in FILES: fakestat(FILES[name])

        elif (CPU_WORD_SIZE==8 and syscall=='fstat' or\
              CPU_WORD_SIZE==4 and syscall=='fstat64') and r[0] in HANDLES:
          fakestat(HANDLES[r[0]]['file'])

        elif syscall=='read' and r[0] in HANDLES:
          H = HANDLES[r[0]]
          data = H['file'][H['pos']:H['pos']+r[2]]
          H['pos'] += len(data)
          write_bytes(pid, r[1], data)
          skipsyscall(len(data))

        elif syscall=='lseek' and r[0] in HANDLES:
          skipsyscall(seek(r[0], r[1], r[2]))

        elif syscall=='_llseek' and r[0] in HANDLES:
          pos = seek(r[0], (r[1]<<32) + designed(r[2]), r[3])
          write_bytes(pid, r[4], struct.pack('=L', pos>>32))
          skipsyscall(pos % 2**32)

        elif syscall=='access' and p(0) in FILES:
          skipsyscall(0)

        elif syscall=='close' and r[0] in HANDLES:
          #del HANDLES[r[0]]
          skipsyscall(0)

        elif syscall=='fcntl' and r[0] in HANDLES:
          if r[1] in (1,3): # F_GETFD F_GETFL
            skipsyscall(1)
          elif r[1] in (2,4): # F_SETFD F_SETFL
            skipsyscall(0)
          else:
            assert False, 'not supported yet %s(...,%d,...)' % (syscall, r[1])

        elif syscall in ('poll', 'select', 'umask', 'fdatasync', 'fsync', 'fstatfs', 'statfs'):
          pass

        # not supported ops
        if syscall in ('ioctl', 'fchmod', 'ftruncate', 'fchown', 'fchdir') and r[0] in HANDLES or\
           syscall in ('creat', 'readlink', 'chown', 'lchown', 'chmod', 'unlink', 'rmdir', 'mkdir', 'truncate', 'mknod') and p(0) in FILES or\
           syscall in ('link', 'symlink', 'rename') and (p(0) in FILES or p(1) in FILES) or\
           syscall in ('unlinkat', 'utimensat') and p(1) in FILES: # TODO: not correct, dirfd to be checked
          skipsyscall(-EINVAL)

        elif syscall=='uselib':
          if DBG: colorwrite( p(0), YELLOW)
          assert p(0) not in FILES, 'not supported yet %s' % syscall

        elif syscall in ('pread64', 'dup', 'dup2', 'flock', 'readv', 'pwrite64', 'write', 'writev'):
          assert r[0] not in HANDLES, 'not supported yet %s' % syscall

        elif syscall in('mmap', 'mmap2'):
          if r[4] in HANDLES:
            r[2], r[3], r[4], r[5] = 7, r[3]|0x20, 0xFFFFFFFF, 0 # ensure PROT_WRITE, MAP_ANONYMOUS
            ptrace(SETREGS, pid, 0, addressof(r))
            if DBG: colorwrite( "new mmap(%x,%x,%x,%x,%x,%x)" % r[0:6], CYAN )

      else: # === on after ===
        ret = P[pid].get('return')
        if ret: # forge return value
          if DBG: colorprint( ' = (%x)' % (ret.acc), RED )
          ptrace(SETREGS, pid, 0, addressof(ret))

        elif P[pid]['args'] is True:
          pass
        else:
          # start, length, prot, flags, fd, offset
          r, syscall = P[pid]['args'], P[pid]['syscall']
          rafter = ptrace_registers_t()
          ptrace(GETREGS, pid, 0, addressof(rafter))

          if syscall in ('mmap','mmap2') and r[4] in HANDLES and not -200<rafter.acc<0:
            start, length, prot, flags, fd, offset = r
            if syscall=='mmap2': offset *= 4096
            write_bytes(pid, rafter.acc, HANDLES[fd]['file'][offset:offset+length])

          if DBG: colorprint( ' = %x' % (rafter.acc), MAGENTA )
        del P[pid]['args']

      ptrace(SYSCALL, pid)
  else:
    ptrace(TRACEME)
    os.execv("/usr/bin/python", ["/usr/bin/python"] + sys.argv + ["--magic"])
    #os.execv("/bin/ls", ["/bin/ls", "-l", "1.so"])

#
print 'kisakuku'

# read file from fakefs
print file('1.so').read()

# visibility of the fake files for child processes
os.system("ls -l 1.py")

# load .so from fakefs
from distorm import Decode, Decode32Bits
print Decode(0x400000, "b800000000".decode("hex"), Decode32Bits)

# execute python script from fakefs
os.system("python 1.py")

# execute python script loading .so from fakefs
os.system("python 2.py")
